

import matplotlib.pyplot as plt  # 补充导入可视化库
from torchvision import transforms as T  # 补充导入transforms
from torchvision.models import resnet50
from modelscope.msdatasets import MsDataset
from modelscope.utils.constant import DownloadMode
from torch.utils.data import DataLoader, Dataset  # 补充导入Dataset
import cv2

# 1. 定义图像预处理
transformer = T.Compose([
    T.Resize((256, 256), antialias=True),
    T.ToTensor()
])


# 2. 自定义数据集类：将ModelScope的字典格式转为 (图像, 标签) 元组
from PIL import Image
import os


class CarDataset(Dataset):
    def __init__(self, ms_dataset, transformer, dataset_name='carBrands50', namespace='tany0699'):
        self.ms_dataset = ms_dataset
        self.transformer = transformer
        # 手动构建数据集根目录（根据ModelScope缓存规则）
        self.data_root = os.path.join(
            os.path.expanduser("~"),  # 获取用户主目录
            '.cache', 'modelscope', 'datasets',
            f'{namespace}%2F{dataset_name}',  # 注意中间是%2F（URL编码的/）
            'default'  # 子集名称，通常是default
        )

    def __len__(self):
        return len(self.ms_dataset)

    def __getitem__(self, idx):
        sample = self.ms_dataset[idx]

        # 1. 获取图像相对路径（替换为实际键名）
        img_path = sample['image:FILE']

        # 2. 拼接完整路径
        full_img_path = os.path.join(self.data_root, img_path)

        # 3. 加载图像
        img = Image.open(full_img_path).convert('RGB')

        # cv2.resize(img, (32, 32),interpolation=cv2.INTER_LINEAR)
        # 4. 预处理和标签提取
        img = self.transformer(img)
        label = int(sample['category'])  # 替换为实际键名

        return img, label


# 3. 加载ModelScope数据集
ms_train_dataset = MsDataset.load(
    'carBrands50',
    namespace='tany0699',
    subset_name='default',
    split='train'
)
ms_val_dataset = MsDataset.load(
    'carBrands50',
    namespace='tany0699',
    subset_name='default',
    split='validation'
)


# 4. 转换为自定义数据集（带预处理）
train_dataset = CarDataset(ms_train_dataset, transformer)
val_dataset = CarDataset(ms_val_dataset, transformer)


# 5. 构建DataLoader
train_dl = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dl = DataLoader(val_dataset, batch_size=45, shuffle=True)


# 6. 可视化训练集第一批的前6张图
if __name__ == "__main__":
    examples = enumerate(train_dl)
    batch_idx, (imgs, labels) = next(examples)
    print(f"当前批次为：{batch_idx}, 图片张量形状为：{imgs.shape}")  # 输出：(128, 3, 32, 32)

    # 创建2行3列的子图
    plt.figure(figsize=(10, 6))  # 设置图的大小
    for i in range(6):
        plt.subplot(2, 3, i+1)
        # 将 (C, H, W) 转为 (H, W, C)，用于plt显示
        img_show = imgs[i].permute(1, 2, 0)
        plt.imshow(img_show)  # 显示RGB图，无需cmap="gray"
        plt.title(f"Label: {labels[i]}")  # 显示标签
        plt.axis('off')  # 隐藏坐标轴，更美观
    plt.tight_layout()  # 自动调整子图间距
    plt.show()

